"""
model.py
返回一个「预训练的 ResNet18」并把最后全连接层改成 5 类输出
"""
import torchvision.models as models  # 官方模型库
import torch.nn as nn                # 网络层
from src.config import Config        # 读取类别数

def get_model():
    if Config.model_name == 'resnet18':
        # 加载在 ImageNet 上预训练的 ResNet18
        net = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

        # 替换最后的全连接层：原来 1000 类 -> 5 类
        # net.fc 是 ResNet18 最后的 nn.Linear 层
        in_features = net.fc.in_features  # 拿到输入维度
        net.fc = nn.Linear(in_features, Config.num_classes)
    else:
        raise NotImplementedError(f'{Config.model_name} 还没实现，请检查 config.py')
    return net